package tjs.study.notes.elasticsearch.select;

import com.alibaba.fastjson.JSONObject;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.Aggregations;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.elasticsearch.search.aggregations.metrics.TopHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import tjs.study.notes.elasticsearch.ElasticsearchApplication;

import java.io.IOException;
import java.util.Arrays;
import java.util.HashMap;
import java.util.stream.Collectors;

/**
 * Created by Nicky.Tang on 2022/8/29 下午5:08
 *
 * @since 02.12.10
 */
@SpringBootTest
public class Agg {
    public static final String groupSplit = "@";
    public static final String groupKey = "groupKey";

    @Autowired
    private RestHighLevelClient restHighLevelClient;

    @Test
    public void base() throws IOException {
        BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery();

        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder()
                .query(boolQueryBuilder)
                .size(11);// 返回多少条明细(分组前)
        searchSourceBuilder.aggregation(AggregationBuilders.terms(groupKey).field("studentName.keyword")
                .size(2)// 返回多少条group by的key
                .subAggregation(AggregationBuilders.topHits("studentNameHits").fetchSource(new String[]{"studentName"}, null)
                        .size(3)));// 返回多少条明细(当前分组)

        SearchRequest searchRequest = new SearchRequest(ElasticsearchApplication.INDEX_NAME);
        searchRequest.source(searchSourceBuilder);

        print(searchRequest);
    }
    
    @Test
    public void bySctipt() throws IOException {
        SearchRequest searchRequest = new SearchRequest(ElasticsearchApplication.INDEX_NAME);
        BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();

        Script script = new Script(String.format("String res = '';for(item in params._source.course){ " +
                        "if(item.courseName!='%s'){continue;}for(itemIns in item.units){res=itemIns.unitName+','+res;} }return res;",
                "语文"));
        searchRequest.source(new SearchSourceBuilder().query(queryBuilder).size(0)
                .aggregation(AggregationBuilders.terms(groupKey).script(script).size(200)));
        print(searchRequest);
    }
    
    @Test
    public void byArray() throws IOException {
        SearchRequest searchRequest = new SearchRequest(ElasticsearchApplication.INDEX_NAME);
        BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();
        searchRequest.source(new SearchSourceBuilder().query(queryBuilder).size(0)
                .aggregation(AggregationBuilders.nested(groupKey, "course")// 根据数组对象分组
                        .subAggregation(buildTermsAggregationBuilder("course.courseName").size(200))));
        print(searchRequest);
    }
    @Test
    public void byArrayFilter() throws IOException {
        SearchRequest searchRequest = new SearchRequest(ElasticsearchApplication.INDEX_NAME);
        BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();
        searchRequest.source(new SearchSourceBuilder().query(queryBuilder).size(0)
                .aggregation(AggregationBuilders.nested(groupKey, "course")
                        .subAggregation(// 过滤数组后再聚合
                                AggregationBuilders.filter("filter", QueryBuilders.termQuery("ucourse.nits", 1))
                        ).subAggregation(
                                AggregationBuilders.terms("a").field("courseName")
                        )
                ));
        print(searchRequest);
    }

    @Test
    public void baseMultiple() throws IOException {
        SearchRequest searchRequest = new SearchRequest(ElasticsearchApplication.INDEX_NAME);
        BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();
        searchRequest.source(new SearchSourceBuilder().query(queryBuilder).size(0)
                .aggregation(buildTermsAggregationBuilder("studentName.keyword", "studentName2.keyword")));
        print(searchRequest);
    }

    public static TermsAggregationBuilder buildTermsAggregationBuilder(String... aggregationFields) {
        String content = Arrays.asList(aggregationFields).stream().map(one -> String.format("doc['%s'].value", one))
                .collect(Collectors.joining("+'" + groupSplit + "'+"));
        TermsAggregationBuilder builder = AggregationBuilders.terms(groupKey);
        builder.script(new Script(ScriptType.INLINE, "painless", content, new HashMap<String, Object>()));
        builder.size(9999);
        return builder;
    }

    private void print(SearchRequest searchRequest) throws IOException {
        SearchResponse searchResponse = restHighLevelClient.search(searchRequest, RequestOptions.DEFAULT);
        Terms terms = searchResponse.getAggregations().get(groupKey);
        for (Terms.Bucket bucket : terms.getBuckets()) {
            System.out.println(bucket.getKeyAsString() + "：：" + bucket.getDocCount() + ">>>>>>>>>>>>>>>>");// 分组key：：count
            Aggregations aggregations = bucket.getAggregations();
            TopHits studentNameHits = aggregations.get("studentNameHits");
            if (studentNameHits == null) {
                continue;
            }
            for (SearchHit hit : studentNameHits.getHits()) {
                System.out.println(JSONObject.toJSONString(hit.getSourceAsMap()));//返回多少条明细(当前分组)
            }
        }
    }

}
